import os.path
import pickle

from tqdm import tqdm
import numpy as np

state_path = './datas/train_state.txt'
text_path = './datas/train_text.txt'
train_path = './datas/train_matrix.pkl'
text_line_num = 0


def get_data_state(data):
    length = len(data)
    if length == 1:
        return "S"
    return "B" + "M" * (length - 2) + "E"


def text_to_state():
    global text_line_num
    if os.path.isfile(state_path):
        return
    all_data = open(text_path, 'r', encoding='utf-8').read().split('\n')
    with open(state_path, 'w', encoding='utf-8') as f:
        length = len(all_data)
        for index, data in tqdm(enumerate(all_data)):
            if not data:
                continue
            text_line_num += 1
            data_splits = data.split(" ")
            tmp_state = ""
            for data_split in data_splits:
                if not data_split:
                    continue
                tmp_state += get_data_state(data_split) + " "
            tmp_state = tmp_state.strip()
            if index != length - 1:
                tmp_state = tmp_state + '\n'
            f.write(tmp_state)


class HMM:
    def __init__(self):

        tmp_texts = open(text_path, 'r', encoding='utf-8').read().split('\n')
        self.all_texts = []
        for ts in tmp_texts:
            if ts:
                tmp_t = ts.split(" ")
                tmp_result = []
                for t in tmp_t:
                    if t:
                        tmp_result.append(t)
                self.all_texts.append(tmp_result)

        # self.all_texts = open(text_path, 'r', encoding='utf-8').read().split('\n')

        tmp_all_states = open(state_path, 'r', encoding='utf-8').read().split('\n')
        self.all_states = []
        for tas in tmp_all_states:
            tas = tas.split(' ')
            self.all_states.append(tas)
        self.states_to_num = {"B": 0, "M": 1, "S": 2, "E": 3}
        self.num_to_states = ["B", "M", "S", "E"]
        self.states_number = len(self.num_to_states)

        self.init_matrix = np.zeros(self.states_number)
        self.transfer_matrix = np.zeros((self.states_number, self.states_number))
        self.emit_matrix = {"B": {"total": 0}, "M": {"total": 0}, "S": {"total": 0}, "E": {"total": 0}}

    def cal_init_matrix(self):
        for state in self.all_states:
            self.init_matrix[self.states_to_num[state[0][0]]] += 1

    def cal_transfer_matrix(self):
        for state in self.all_states:
            state_join = "".join(state)
            state_pre = state_join[:-1]
            state_nxt = state_join[1:]
            for sp, sn in zip(state_pre, state_nxt):
                self.transfer_matrix[self.states_to_num[sp], self.states_to_num[sn]] += 1

    def cal_emit_matrix(self):
        for one_text, one_state in zip(self.all_texts, self.all_states):
            for text, state in zip("".join(one_text), "".join(one_state)):

                if text in self.emit_matrix[state]:
                    self.emit_matrix[state][text] += 1
                else:
                    self.emit_matrix[state][text] = 1

                # self.emit_matrix[state][word] = self.emit_matrix[state].get(word, 0) + 1

                self.emit_matrix[state]["total"] += 1

    def normalize(self):
        self.init_matrix = self.init_matrix / np.sum(self.init_matrix)
        self.transfer_matrix = self.transfer_matrix / np.sum(self.transfer_matrix, axis=1, keepdims=True)
        for state, word_detail in self.emit_matrix.items():
            for word, num in word_detail.items():
                if word == 'total':
                    continue
                self.emit_matrix[state][word] = num / self.emit_matrix[state]['total']

    def train(self):
        global train_path
        if os.path.isfile(train_path):
            self.init_matrix, self.transfer_matrix, self.emit_matrix = pickle.load(open(train_path, "rb"))
            # print(self.init_matrix)
            # print(self.transfer_matrix)
            # print(self.emit_matrix)
            return
        self.cal_init_matrix()
        self.cal_transfer_matrix()
        self.cal_emit_matrix()
        self.normalize()
        pickle.dump([self.init_matrix, self.transfer_matrix, self.emit_matrix], open(train_path, "wb"))


def viterbi_t(text, hmm):
    length = len(text)
    if length == 0:
        return
    init_p = hmm.init_matrix
    states = hmm.num_to_states
    emit_p = hmm.emit_matrix
    number = hmm.states_to_num
    transfer_p = hmm.transfer_matrix
    tmp_result = [1, 1, 1, 1]
    path = []
    for i in range(length):
        tmp_arr = []
        for j in range(4):
            tmp_arr.append('S')
        path.append(tmp_arr)
    for index, state in enumerate(states):
        tmp_result[index] = tmp_result[index] * init_p[index] * emit_p[state].get(text[0], 0)
        path[0][index] = state
    for i in range(1, length):
        word = text[i]
        neverSeen = word not in emit_p['S'].keys() and word not in emit_p['B'].keys() and word not in emit_p['E'].keys() and word not in emit_p['M'].keys()
        now_result = [0, 0, 0, 0]
        for index, state in enumerate(states):
            if neverSeen:
                nowEmitP = 1
            else:
                nowEmitP = emit_p[state].get(word, 0)
            for i1, s1 in enumerate(states):
                if tmp_result[i1] * nowEmitP * transfer_p[i1][index] >= now_result[index]:
                    path[i][index] = s1
                    now_result[index] = tmp_result[i1] * nowEmitP * transfer_p[i1][index]
        for sn in range(hmm.states_number):
            tmp_result[sn] = now_result[sn]
    max_ans = 0
    pos = ''
    start_ind = length - 1
    ans = []
    for i in range(hmm.states_number):
        if tmp_result[i] > max_ans:
            max_ans = tmp_result[i]
            pos = states[i]
    ans.append(pos)
    while start_ind > 0:
        fr = path[start_ind][number[pos]]
        start_ind -= 1
        ans.append(fr)
        pos = fr
    ans = reversed(ans)
    final_ans = get_final_ans(text, ans)


def get_final_ans(text, ans):
    final_ans = ""
    for t, a in zip(text, ans):
        if a == 'S' or a == 'E':
            final_ans += t + ' '
        elif a == 'B' or a == 'M':
            final_ans += t
    return final_ans


text_to_state()
hmm = HMM()
hmm.train()
while True:
    sen = input()
    viterbi_t(sen, hmm)

